package com.LeeCode;

import java.util.*;

/**
 * 连接两棵树后最大目标节点数目 I
 */

public class Code3372 {
    public static void main(String[] args) {
        int[][] edges1 = {{0, 1}, {0, 2}, {2, 3}, {2, 4}};
        int[][] edges2 = {{0, 1}, {0, 2}, {0, 3}, {2, 7}, {1, 4}, {4, 5}, {4, 6}};
        int k = 2;

        System.out.println(Arrays.toString(new Code3372().maxTargetNodes(edges1, edges2, k)));
    }

    public int[] maxTargetNodes(int[][] edges1, int[][] edges2, int k) {
        int n = edges1.length + 1, m = edges2.length + 1;
        int[] ans = new int[n];
        HashMap<Integer, List<Integer>> edgeMap1 = new HashMap<>();
        HashMap<Integer, List<Integer>> edgeMap2 = new HashMap<>();

        for (int[] edge : edges1) {
            if (!edgeMap1.containsKey(edge[0]))
                edgeMap1.put(edge[0], new ArrayList<>());
            if (!edgeMap1.containsKey(edge[1]))
                edgeMap1.put(edge[1], new ArrayList<>());

            edgeMap1.get(edge[0]).add(edge[1]);
            edgeMap1.get(edge[1]).add(edge[0]);
        }

        for (int[] edge : edges2) {
            if (!edgeMap2.containsKey(edge[0]))
                edgeMap2.put(edge[0], new ArrayList<>());
            if (!edgeMap2.containsKey(edge[1]))
                edgeMap2.put(edge[1], new ArrayList<>());
            edgeMap2.get(edge[0]).add(edge[1]);
            edgeMap2.get(edge[1]).add(edge[0]);
        }

        int max2 = 0;
        if (k > 0) {
            for (int i = 0; i < m; i++) {
                max2 = Math.max(max2, bfs(edgeMap2, new int[m], i, k - 1));
            }
        }

        for (int i = 0; i < n; i++) {
            ans[i] = bfs(edgeMap1, new int[n], i, k) + max2;
        }
        return ans;
    }

    private int bfs(HashMap<Integer, List<Integer>> edgeMap, int[] visited, int start, int k) {
        int ans = 0;
        Queue<Integer> queue = new LinkedList<>();
        queue.add(start);
        int turn = 0;

        while (!queue.isEmpty() && turn <= k) {
            int size = queue.size();
            for (int i = 0; i < size; i++) {
                int cur = queue.poll();
                if (visited[cur] == 1)
                    continue;
                visited[cur] = 1;
                if (edgeMap.containsKey(cur)) {
                    queue.addAll(edgeMap.get(cur));
                }
                ans++;
            }
            turn++;
        }
        return ans;
    }
}

